# -*- coding: utf-8 -*-
'''
@File    :   agentApi.py
@Author  :   一力辉 
'''

from .reponse import BaseResponse, Response
from typing import Optional, List
from pydantic import BaseModel
from fastapi import APIRouter, Request
from fastapi.responses import JSONResponse, StreamingResponse
from digitalHuman.utils import config
from digitalHuman.utils import TextMessage
from digitalHuman.agent import AgentPool
from sse_starlette.sse import EventSourceResponse
from dashscope import Application
from digitalHuman.utils import logger
import re
import httpx

import json
import asyncio  

router = APIRouter()
agentPool = AgentPool()

class InferIn(BaseModel):
    engine: str = "default"
    settings: dict = {}
    streaming: bool = False
    data: str

class InferOut(BaseResponse):
    data: bytes

async def interalError():
    yield "内部错误"

# @router.get("/v0/infer", response_model=InferOut, summary="AI Agent Inference")
@router.get("/v0/infer")
async def apiInfer(prompt: str, req: Request):
    # if item.engine.lower() == "default":
    #     item.engine = config.SERVER.AGENTS.DEFAULT
    # item = InferIn()
    # item.engine = config.SERVER.AGENTS.DEFAULT
    # item.streaming = True
    response = Response()
    
    try:
        input = TextMessage(data=prompt)
        # dify服务参数校验
        # if "dify" in item.engine.lower():
        #     if "url" not in item.settings or "key" not in item.settings:
        #         raise RuntimeError("dify url and key is required")
        
        # responses = Application.call(app_id='6d3a691a572447b1bd91b7fc23d3c54a',
        #                     prompt='京杭大运河分哪几个段？',
        #                     api_key='sk-32ea4749231641759586bbd473522e8b',
        #                     )
        
        # async def event_generator(request: Request):  
        #     for response in responses:  
        #         if await request.is_disconnected():  
        #             print("连接已中断")  
        #             break  
        #         yield bytes(response.output.text, encoding='utf-8')

        # return EventSourceResponse(event_generator(req))

        # return StreamingResponse(agentPool.get(item.engine).run(input, item.streaming, **item.settings))
        # result = agentPool.get(item.engine).run(input, item.streaming, **item.settings)
        # return StreamingResponse(result)
        # return StreamingResponse(result, media_type="text/event-stream")
        # return EventSourceResponse(agentPool.get(item.engine).run(input, item.streaming, **item.settings))
        logger.info(f"开始调用提问:{config.SERVER.AGENTS.DEFAULT}，提问文字:{prompt}")
        result = agentPool.get(config.SERVER.AGENTS.DEFAULT).run(input, True)
        logger.info(f"调用提问结束")
        return EventSourceResponse(result)
    
    except Exception as e:
        response.error(str(e))
        return EventSourceResponse(interalError)


class ListOut(BaseResponse):
    data: List[str] = []

@router.get("/v0/list", response_model=ListOut, summary="AI Agent List")
async def apiList():
    response = Response()
    try:
        response.data = agentPool.list()
    except Exception as e:
        response.error(str(e))
    return JSONResponse(content=response.validate(ListOut), status_code=200)

class Message(BaseModel):  
    message: str  

@router.get("/v0/sse")  
async def indexpost(prompt: str, req: Request):  
    async def event_generator(request: Request):  
        KEY = "sk-32ea4749231641759586bbd473522e8b"
        APP_ID = "6d3a691a572447b1bd91b7fc23d3c54a"
        headers = {
                'Content-Type': 'application/json',
                'Authorization': f'Bearer {KEY}',
                'X-DashScope-SSE': 'enable'
            }
        payload = {
            "input": {"prompt": "介绍一下自己"},
            "parameters":  {"incremental_output": "true"},
            "debug": {}
        }

        pattern = re.compile(r'data:\s*({.*})')
        client = httpx.AsyncClient(headers=headers)

        async with client.stream('POST', f'https://dashscope.aliyuncs.com/api/v1/apps/{APP_ID}/completion', headers=headers, json=payload) as response:
            async for chunk in response.aiter_bytes():
                if await request.is_disconnected():  
                    print("连接已中断")  
                    break  
                chunkStr = chunk.decode('utf-8').strip()
                # print(chunkStr)
                chunkData = pattern.search(chunkStr)
                # 部分dify返回不完整，该模板匹配会失效
                if not chunkData: 
                    print(f"[AGENT] Engine return truncated data: {chunkData}")
                    continue
                chunkData = chunkData.group(1)

                try:
                    data = json.loads(chunkData)
                    if data['output']['finish_reason'] == "stop":
                        logger.info(f"AliAgent输出结束。")
                        yield {"data" : "", "finish_status" : "true"}
                    else:
                        logger.info(f"AliAgent输出： {data['output']['text']}")
                        yield {"data" : data['output']['text'], "finish_status" : "false"}
                except Exception as e:
                    logger.info(f"[AGENT] Engine run failed: {e}")
                    yield {"data" : "内部错误，请检查dify信息。" , "finish_status" : "true"}

    return EventSourceResponse(event_generator(req))
